import os
import cv2
from tqdm import tqdm

def draw_yolo_boxes(image, label_path, color=(0, 255, 0), thickness=2):
    """
    Draw YOLO format bounding boxes on an image.
    """
    h, w = image.shape[:2]
    if not os.path.exists(label_path):
        return image  # If no labels, return the original image

    with open(label_path, 'r') as f:
        lines = f.readlines()

    for line in lines:
        parts = line.strip().split()
        if len(parts) != 5:
            continue
        cls_id, x_center, y_center, bbox_w, bbox_h = map(float, parts)
        x_center *= w
        y_center *= h
        bbox_w *= w
        bbox_h *= h

        x1 = int(x_center - bbox_w / 2)
        y1 = int(y_center - bbox_h / 2)
        x2 = int(x_center + bbox_w / 2)
        y2 = int(y_center + bbox_h / 2)

        cv2.rectangle(image, (x1, y1), (x2, y2), color, thickness)
        cv2.putText(image, str(int(cls_id)), (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX,
                    0.5, color, 1, lineType=cv2.LINE_AA)

    return image

def visualize_yolo_results(image_dir, label_dir, vis_output_dir):
    """
    Visualizes YOLO results by drawing bounding boxes on split images.
    """
    os.makedirs(vis_output_dir, exist_ok=True)
    image_files = [f for f in os.listdir(image_dir) if f.lower().endswith(('.jpg', '.png'))]

    for img_file in tqdm(image_files, desc="Visualizing"):
        image_path = os.path.join(image_dir, img_file)
        label_path = os.path.join(label_dir, img_file.replace('.jpg', '.txt').replace('.png', '.txt'))
        image = cv2.imread(image_path,cv2.IMREAD_UNCHANGED)
        if image is None:
            print(f"Warning: failed to read {image_path}")
            continue

        vis_image = draw_yolo_boxes(image, label_path)
        vis_path = os.path.join(vis_output_dir, img_file)
        cv2.imwrite(vis_path, vis_image)

if __name__ == "__main__":
    image_dir = "/home/lhx/data_zoo/splittruecolorwithdepthseg/images"
    label_dir = "/home/lhx/data_zoo/splittruecolorwithdepthseg/labels"
    vis_output_dir = "/home/lhx/data_zoo/splittruecolorwithdepthseg/vis"
    visualize_yolo_results(image_dir, label_dir, vis_output_dir)
